import torch
from tqdm import tqdm

from config import config
from data.dataset import get_data_loaders
from models.cnn import CIFAR10CNN
from utils.logger import setup_logger

from pathlib import Path

current_file = Path(__file__).name
logger = setup_logger(current_file)


def test_model():
    # 获取数据加载器
    _, test_loader, classes = get_data_loaders()

    # 初始化模型
    model = CIFAR10CNN(num_classes=len(classes)).to(config.DEVICE)

    # 加载训练好的模型
    try:
        checkpoint = torch.load(config.BEST_MODEL_PATH)
        model.load_state_dict(checkpoint['model_state_dict'])
        model.eval()
        logger.info("Model loaded successfully")
    except FileNotFoundError:
        logger.error(f"No trained model found at {config.BEST_MODEL_PATH}")
        return

    # 测试模型
    correct = 0
    total = 0

    with torch.no_grad():
        for images, labels in tqdm(test_loader, desc="Testing"):
            images, labels = images.to(config.DEVICE), labels.to(config.DEVICE)
            outputs = model(images)
            _, predicted = torch.max(outputs.data, 1)
            total += labels.size(0)
            correct += (predicted == labels).sum().item()

    accuracy = 100 * correct / total
    logger.info(f"Test Accuracy: {accuracy:.2f}%")

    # 打印每个类别的准确率
    class_correct = [0] * len(classes)
    class_total = [0] * len(classes)

    with torch.no_grad():
        for images, labels in test_loader:
            images, labels = images.to(config.DEVICE), labels.to(config.DEVICE)
            outputs = model(images)
            _, predicted = torch.max(outputs, 1)
            c = (predicted == labels).squeeze()

            for i in range(len(labels)):
                label = labels[i]
                class_correct[label] += c[i].item()
                class_total[label] += 1

    for i in range(len(classes)):
        logger.info(
            f"Accuracy of {classes[i]:10s}: {100 * class_correct[i] / class_total[i]:.2f}% "
            f"({class_correct[i]}/{class_total[i]})"
        )


if __name__ == '__main__':
    test_model()
